#include "kernel_operator.h"
using namespace AscendC;

class KernelCumsum
{
public:
    __aicore__ inline KernelCumsum()
    {
    }
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR axis, GM_ADDR y, int64_t blockLength,
                                int64_t shape[], uint32_t dim_num, int32_t exclusive, int32_t reverse)
    {
        xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x, blockLength);
    axisGm.SetGlobalBuffer((__gm__ DTYPE_AXIS *)axis, 1);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y, blockLength);
    int32_t axisScalar = axisGm.GetValue(0);
    if (axisScalar < 0){
        axisScalar += dim_num;
    }
    int32_t length = shape[axisScalar];
    int32_t rpttime = blockLength/length;
    int32_t stepsize = 1;
    for (int i = axisScalar + 1; i < dim_num; i++)
    {
        stepsize = stepsize * shape[i];
    }

    if (reverse)
    {
        if (exclusive)
        {
            for (int outrpt = 0; outrpt < rpttime; outrpt++)
            {
                for (int inrpt = 0; inrpt < length; inrpt++)
                {
                    if (inrpt == 0)
                    {
                        yGm.SetValue(((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + (length - inrpt - 1) * stepsize), (float)0);
                    }
                    else
                    {
                        float Tvalue = (float)yGm.GetValue((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + (length - inrpt) * stepsize) + (float)xGm.GetValue((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + (length - inrpt) * stepsize);
                        yGm.SetValue(((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + (length - inrpt - 1) * stepsize), (DTYPE_X)Tvalue);
                    }
                }
            }
        }else
        {
            for (int outrpt = 0; outrpt < rpttime; outrpt++)
            {
                for (int inrpt = 0; inrpt < length; inrpt++)
                {
                    if (inrpt == 0)
                    {
                        float Tvalue = (float)xGm.GetValue((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + (length - inrpt - 1) * stepsize);
                        yGm.SetValue(((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + (length - inrpt - 1) * stepsize), (DTYPE_X)Tvalue);
                    }
                    else
                    {
                        float Tvalue = (float)yGm.GetValue((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + (length - inrpt) * stepsize) + (float)xGm.GetValue((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + (length - inrpt - 1) * stepsize);
                        yGm.SetValue(((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + (length - inrpt - 1) * stepsize), (DTYPE_X)Tvalue);
                    }
                }
            }
        }
    }else
    {
        if (exclusive)
        {
            for (int outrpt = 0; outrpt < rpttime; outrpt++)
            {
                for (int inrpt = 0; inrpt < length; inrpt++)
                {
                    if (inrpt == 0)
                    {
                        yGm.SetValue(((outrpt / stepsize) * length * stepsize + (outrpt % stepsize)), (DTYPE_X)0);
                    }
                    else
                    {
                        float Tvalue = (float)yGm.GetValue((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + (inrpt - 1) * stepsize) + (float)xGm.GetValue((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + (inrpt - 1) * stepsize);
                        yGm.SetValue(((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + inrpt * stepsize), (DTYPE_X)Tvalue);
                    }
                }
            }
        } else
        {
            for (int outrpt = 0; outrpt < rpttime; outrpt++)
            {
                for (int inrpt = 0; inrpt < length; inrpt++)
                {
                    if (inrpt == 0)
                    {
                        float Tvalue = xGm.GetValue((outrpt / stepsize) * length * stepsize + (outrpt % stepsize));
                        yGm.SetValue(((outrpt / stepsize) * length * stepsize + (outrpt % stepsize)), (DTYPE_X)Tvalue);
                    }
                    else
                    {
                        float Tvalue = (float)xGm.GetValue((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + inrpt * stepsize) + (float)yGm.GetValue((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + (inrpt - 1) * stepsize);
                        yGm.SetValue(((outrpt / stepsize) * length * stepsize + (outrpt % stepsize) + inrpt * stepsize), (DTYPE_X)Tvalue);
                    }
                }
            }
        }
    }
}

private:
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_AXIS> axisGm;
    GlobalTensor<DTYPE_Y> yGm;

};
extern "C" __global__ __aicore__ void cumsum(GM_ADDR x, GM_ADDR axis, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelCumsum op;
    op.Init(x, axis, y, tiling_data.blockLength,
            tiling_data.shape, tiling_data.dimnum, tiling_data.exclusive, tiling_data.reverse);
}


